import torch
import torchvision
import torch.nn as nn
import torch.utils.data as td
import matplotlib.pyplot as plt
from torch.autograd import Variable
from cnn import CNN

num_epochs = 1  # 训练轮数
batch_size = 50  # 每个训练批次所包含的样本数量
learning_rate = 0.001  # 学习率
download_mnist = False  # 是否下载 MNIST 数据集

train_data = torchvision.datasets.MNIST(
    root='./',  # 数据集存放的位置
    train=True,  # 加载训练集
    transform=torchvision.transforms.ToTensor(),  # 数据预处理, 将图片转换为 Tensor 并归一化
    download=download_mnist  # 是否下载数据集
)

# 绘制训练数据
# print(train_data.data.size())
# print(train_data.targets.size())
# plt.imshow(train_data.data[1].numpy(), cmap='gray')
# plt.title('%i' % train_data.targets[1])
# plt.show()

# 创建训练数据集的数据加载器
train_loader = td.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True)

test_data = torchvision.datasets.MNIST(
    root='./',
    train=False
)
with torch.no_grad():
    test_x = torch.unsqueeze(test_data.data.type(torch.FloatTensor), dim=1)[:2000] / 255
    test_y = test_data.targets[:2000]

if __name__ == '__main__':
    # 创建模型
    cnn = CNN()
    # 设置优化器与损失函数
    optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
    loss_function = nn.CrossEntropyLoss()
    # 训练模型
    step = 0
    for epoch in range(num_epochs):
        for step, data in enumerate(train_loader):
            x, y = data
            output = cnn(x)  # 前向传播, 计算模型输出
            loss = loss_function(output, y)  # 计算损失
            optimizer.zero_grad()  # 梯度清零
            loss.backward()  # 反向传播, 计算梯度
            optimizer.step()  # 更新模型参数

            # 每执行 50 次, 输出一下当前的 epoch / loss / accuracy
            if step % 50 == 0:
                test_output = cnn(test_x)
                test_y_predicted = torch.max(test_output, 1)[1].data.squeeze()
                accuracy = torch.eq(test_y_predicted, test_y).sum().item() / test_y.size(0)
                print(f'Now epoch: {epoch + 1}, loss: {loss.item()}, accuracy: {accuracy}')

    test_output = cnn(test_x[:50])
    y_pred = torch.max(test_output, 1)[1].data.squeeze()
    print('Predict Result:', y_pred.tolist())
    print('Real Result:', test_y[:50].tolist())
